#!/bin/bash
#===============================================================================
# MPI性能诊断自动化测试脚本
# 功能：自动运行不同核数的测试，收集性能数据并生成报告
#
# 使用方法：
#   ./run_mpi_scaling_test.sh
#
# 说明：
#   - 脚本会自动测试多个核数配置
#   - 生成详细的性能报告
#   - 帮助判断是硬件问题还是算法问题
#===============================================================================

# export LD_PRELOAD=/root/pacific_ext/WUCHAO_users/jiangxu/mpiTest/mpiTimer/libmpitimer.so

# 设置颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 获取脚本所在目录
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "${SCRIPT_DIR}"

# 定义测试的核数数组（根据你的608核服务器调整）
# 建议测试范围：从小到大，观察扩展性曲线
# CORE_COUNTS=(9 18 36 72 144 288 576)
CORE_COUNTS=(128)

# 结果输出目录和文件
RESULT_DIR="${SCRIPT_DIR}/results"
# TIMESTAMP=$(date +%Y%m%d_%H%M%S)
RESULT_FILE="${RESULT_DIR}/mpi_performance.txt"
SUMMARY_FILE="${RESULT_DIR}/scaling_summary.csv"

# 创建结果目录
mkdir -p "${RESULT_DIR}"

#===============================================================================
# 函数定义
#===============================================================================

# 打印带颜色的消息
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查可执行文件
check_executable() {
    if [ ! -f "${SCRIPT_DIR}/mpiTest" ]; then
        print_error "未找到可执行文件 mpiTest"
        print_info "请先运行编译脚本: ./compile.sh"
        exit 1
    fi
}

# 运行单次测试
run_single_test() {
    local nproc=$1
    local log_file="${RESULT_DIR}/mpi_test_${nproc}cores.log"
    
    print_info "=========================================="
    print_info "测试配置: ${nproc} 核"
    print_info "=========================================="
    
    # MPI运行参数（针对鲲鹏优化）
    # --bind-to core: 将进程绑定到核心
    # --map-by core: 按核心映射
    # UCX参数：优化网络通信
    local mpi_params=""
    
    # 如果使用OpenMPI
    if mpirun --version 2>&1 | grep -q "Open MPI"; then
        mpi_params="--bind-to core --map-by core"
        # 对于鲲鹏，可能需要特殊的UCX设置
        export UCX_RC_VERBS_TX_MIN_SGE=2
        export UCX_UD_VERBS_TX_MIN_SGE=1
    fi
    
    # 运行测试
    print_info "开始运行测试..."
    local start_time=$(date +%s)
    
    # mpirun --allow-run-as-root -x UCX_RC_VERBS_TX_MIN_SGE=2 -x UCX_UD_VERBS_TX_MIN_SGE=1 ${mpi_params} -np ${nproc} --bind-to core --map-by numa:PE=1 ./mpiTest > "${log_file}" 2>&1

    mpirun --prefix $MY_MPI_ROOT \
            --allow-run-as-root \
            -np ${nproc} \
            --map-by numa:SPAN \
            --bind-to core \
            --mca coll_han_use_dynamic_file_rules 1 \
            --mca coll_han_dynamic_rules_filename /root/pacific_ext/WUCHAO_users/jiangxu/share/tools/ompi5-improved/tuning-files-aws/han_rules_1 \
            --mca coll_tuned_use_dynamic_rules 1 \
            --mca coll_tuned_dynamic_rules_filename /root/pacific_ext/WUCHAO_users/jiangxu/share/tools/ompi5-improved/tuning-files-aws/tuned_rules_1 \
            ./mpiTest > "${log_file}" 2>&1
    local exit_code=$?
    
    local end_time=$(date +%s)
    local wall_time=$((end_time - start_time))
    
    if [ ${exit_code} -eq 0 ]; then
        print_success "${nproc}核测试完成！用时: ${wall_time}秒"
        print_info "日志文件: ${log_file}"
        
        # 追加到汇总文件
        echo "========================================" >> "${RESULT_FILE}"
        echo "测试配置: ${nproc} 核" >> "${RESULT_FILE}"
        echo "完成时间: $(date)" >> "${RESULT_FILE}"
        echo "总用时: ${wall_time} 秒" >> "${RESULT_FILE}"
        echo "========================================" >> "${RESULT_FILE}"
        cat "${log_file}" >> "${RESULT_FILE}"
        echo "" >> "${RESULT_FILE}"
        
        return 0
    else
        print_error "${nproc}核测试失败！"
        print_error "查看日志: ${log_file}"
        return 1
    fi
}

# 解析结果并生成CSV汇总
generate_summary() {
    print_info "生成性能汇总表..."
    
    # CSV表头
    echo "核数,强扩展计算时间(s),内存带宽(GB/s),负载不均衡度(%),并行效率(%)" > "${SUMMARY_FILE}"
    
    for nproc in "${CORE_COUNTS[@]}"; do
        local log_file="${RESULT_DIR}/mpi_test_${nproc}cores.log"
        
        if [ -f "${log_file}" ]; then
            # 提取关键性能指标
            local compute_time=$(grep "计算时间:" "${log_file}" | tail -1 | awk '{print $2}')
            local mem_bw=$(grep "平均内存带宽:" "${log_file}" | awk '{print $2}')
            local load_imbalance=$(grep "负载不均衡度:" "${log_file}" | awk '{print $2}' | sed 's/%//')
            local parallel_eff=$(grep "并行效率:" "${log_file}" | awk '{print $2}' | sed 's/%//')
            
            # 写入CSV
            echo "${nproc},${compute_time:-N/A},${mem_bw:-N/A},${load_imbalance:-N/A},${parallel_eff:-N/A}" >> "${SUMMARY_FILE}"
        fi
    done
    
    print_success "汇总表已保存: ${SUMMARY_FILE}"
}

# 生成诊断报告
generate_diagnostic_report() {
    local report_file="${RESULT_DIR}/diagnostic_report.txt"
    
    print_info "生成诊断报告..."
    
    {
        echo "================================================================================"
        echo "  MPI并行性能诊断报告"
        echo "================================================================================"
        echo "生成时间: $(date)"
        echo "测试核数范围: ${CORE_COUNTS[@]}"
        echo ""
        echo "================================================================================"
        echo "  性能指标汇总"
        echo "================================================================================"
        echo ""
        
        # 读取CSV并格式化显示
        column -t -s ',' "${SUMMARY_FILE}"
        
        # echo ""
        # echo "================================================================================"
        # echo "  诊断建议"
        # echo "================================================================================"
        # echo ""
        # echo "1. 通信性能分析："
        # echo "   - 检查点对点通信延迟是否随核数增加而上升"
        # echo "   - 如果延迟 > 2us，可能是网络配置问题"
        # echo "   - 如果带宽 < 10 GB/s（对于InfiniBand），检查网络设备"
        # echo ""
        # echo "2. 内存带宽分析："
        # echo "   - 鲲鹏920理论内存带宽: ~150 GB/s per socket"
        # echo "   - 如果实测带宽 < 100 GB/s，可能有NUMA问题"
        # echo "   - 如果各进程带宽差异 > 20%，检查NUMA绑定"
        # echo ""
        # echo "3. 负载均衡分析："
        # echo "   - 负载不均衡度应 < 5%"
        # echo "   - 如果 > 10%，可能是进程绑定或算法问题"
        # echo "   - 检查是否有进程共享CPU核心"
        # echo ""
        # echo "4. 扩展性分析："
        # echo "   - 绘制加速比曲线：加速比 = T(1) / T(n)"
        # echo "   - 理想扩展性：加速比 = n"
        # echo "   - 如果加速比曲线快速下降："
        # echo "     * 可能是通信开销过大（算法问题）"
        # echo "     * 可能是同步开销大（负载不均衡）"
        # echo "     * 可能是内存带宽瓶颈（硬件限制）"
        # echo ""
        # echo "5. 计算密集型任务分析："
        # echo "   - 检查GFLOPS是否随核数线性增长"
        # echo "   - 如果不是，说明可能有以下问题："
        # echo "     * 硬件资源竞争（缓存、内存带宽）"
        # echo "     * 进程绑定不当"
        # echo "     * NUMA配置问题"
        # echo ""
        # echo "================================================================================"
        # echo "  进一步测试建议"
        # echo "================================================================================"
        # echo ""
        # echo "1. 如果怀疑是硬件问题："
        # echo "   - 运行STREAM内存带宽测试"
        # echo "   - 运行HPL (High Performance Linpack)"
        # echo "   - 检查NUMA拓扑: numactl --hardware"
        # echo "   - 检查CPU频率: lscpu | grep MHz"
        # echo ""
        # echo "2. 如果怀疑是算法问题："
        # echo "   - 分析算法的通信-计算比"
        # echo "   - 使用性能分析工具（如TAU, VTune, perf）"
        # echo "   - 检查同步点和全局通信频率"
        # echo "   - 考虑改进域分解策略"
        # echo ""
        # echo "3. 如果怀疑是MPI配置问题："
        # echo "   - 尝试不同的进程绑定策略"
        # echo "   - 调整MPI缓冲区大小"
        # echo "   - 测试不同的MPI实现（OpenMPI vs MPICH）"
        # echo ""
        # echo "================================================================================"
        
    } > "${report_file}"
    
    print_success "诊断报告已生成: ${report_file}"
    
    # 显示报告
    cat "${report_file}"
}

#===============================================================================
# 主程序
#===============================================================================

main() {
    print_info "=========================================="
    print_info "MPI并行性能诊断自动化测试"
    print_info "=========================================="
    echo ""
    
    # 检查可执行文件
    check_executable
    
    # 初始化结果文件
    echo "MPI并行性能诊断测试结果" > "${RESULT_FILE}"
    echo "开始时间: $(date)" >> "${RESULT_FILE}"
    echo "测试核数: ${CORE_COUNTS[@]}" >> "${RESULT_FILE}"
    echo "" >> "${RESULT_FILE}"
    
    # 显示系统信息
    print_info "系统信息:"
    print_info "  - CPU架构: $(uname -m)"
    print_info "  - 操作系统: $(cat /etc/os-release | grep PRETTY_NAME | cut -d'=' -f2 | tr -d '\"')"
    print_info "  - CPU核心数: $(nproc)"
    print_info "  - MPI版本: $(mpirun --version | head -1)"
    echo ""
    
    # 循环运行测试
    local failed_tests=0
    for nproc in "${CORE_COUNTS[@]}"; do
        # 检查核数是否超过可用核心
        local available_cores=$(nproc)
        if [ ${nproc} -gt ${available_cores} ]; then
            print_warning "跳过 ${nproc} 核测试（超过可用核心数 ${available_cores}）"
            continue
        fi
        
        run_single_test ${nproc}
        
        if [ $? -ne 0 ]; then
            ((failed_tests++))
            print_warning "测试失败，继续下一个..."
        fi
        
        echo ""
        sleep 2  # 短暂延迟，让系统恢复
    done
    
    # 生成汇总和报告
    generate_summary
    generate_diagnostic_report
    
    # 输出最终结果
    print_success "=========================================="
    print_success "所有测试完成！"
    print_success "=========================================="
    print_info "完整日志: ${RESULT_FILE}"
    print_info "性能汇总: ${SUMMARY_FILE}"
    print_info "诊断报告: ${RESULT_DIR}/diagnostic_report.txt"
    
    if [ ${failed_tests} -gt 0 ]; then
        print_warning "警告: ${failed_tests} 个测试失败"
    fi
    
    echo ""
    print_info "建议：使用Excel或其他工具打开CSV文件，绘制扩展性曲线"
}

# 运行主程序
main

